#Copyright (C) Nial Peters 2012
#
#This file is part of AvoScan.
#
#AvoScan is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.
#
#AvoScan is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with AvoScan.  If not, see <http://www.gnu.org/licenses/>.
from Phidgets import PhidgetException
"""
The programs module contains code for running different types of scan. 
In general all scan programs should inherit from ScanProgramBase and 
must implement a _run(self) method which will be run in a separate thread
when start() is called. The _run() method should return as soon as _exit_flag
becomes True.
"""

import threading
import datetime
import Queue

from avoscan.processing import date2secs


class ScanProgramBase:
    """
    Base class for scan programs - must be subclassed. All scan programs should inherit from 
    ScanProgramBase and implement a _run() method (which takes no args). This will be run in a 
    separate thread when start() is called. The _run() method should return as soon as
    _exit_flag becomes True.
    
    The scanner argument should be a scanner.Scanner object.
    """
    
    def __init__(self, scanner, start_angle, extra_procs=None):       
        self._scanner = scanner
        self.__scanning_thread = None
        self._exit_flag = False
        self._is_scanning = False
        self._extra_procs = extra_procs
        self.start_angle = float(start_angle)
      
      
    def _run(self):
        """
        Method that actually implements the scan program. This must be overridden 
        by the subclass. It will be run in a separate thread when start() is called.
        It should return as soon as _exit_flag becomes True. See the ScanProgramSimple
        class for an example of how to implement the _run method.
        """
        raise NotImplementedError, "ScanProgramBase must be sub-classed."
        
        
    def stop(self):
        """
        Stops the scan program.
        """
        self._exit_flag = True
        
        #stop the extra threads/processes BEFORE the scanning_thread
        if self._extra_procs is not None:
            for x in self._extra_procs:
                x.stop()
        
        if self.__scanning_thread is not None:
            try:
                self._scanner.stop()
            except PhidgetException.PhidgetException:
                #the stepper might be disconnected
                pass
            
            self.__scanning_thread.join()
        
        self._is_scanning = False
        
    
    
    def start(self):
        """
        Starts the scan program.
        """
        self._exit_flag = False
        self._scanner.engage()
        self._goto_start_pos()
        self.__scanning_thread = threading.Thread(target=self._run)
        self._is_scanning = True
        if self._extra_procs is not None:
            for x in self._extra_procs:
                x.start()
        self.__scanning_thread.start()
    
    
    def _goto_start_pos(self):
        """
        Rotates the scanner to the starting angle. 
        """
        self._scanner.set_to_tracking_speed()
        self._scanner.set_position(self.start_angle)



class ScanProgramSimple(ScanProgramBase):
    """
    Simple scan program that rotates the scanner between the start and end angles
    at the desired velocity.
    
    * scanner - a scanner.Scanner object
    * start - start angle in degrees
    * end - end angle in degrees
    * duration - scan time (one way) in minutes.
    * extra_procs - extra processes that need to be run at the same time as the scan program
      for example, this is how the file system monitor threads are started and ended. It should
      be a list of objects that implement a start and stop method. The start method will be called
      on scan program start and stop will be called on scan program end. The start method must 
      return immediately.
    """
    
    def __init__(self, scanner, start, end, duration, extra_procs=None): 
        ScanProgramBase.__init__(self, scanner, start, extra_procs=extra_procs)
        self.end_angle = float(end)
        self.duration = float(duration)


    def set_start_angle(self, angle):
        """
        Sets the start angle for the program to 'angle', where
        'angle' is in degrees. This method may be called while the 
        scan program is running. However, the start angle listed at 
        the top of the angles file will not be changed.
        """
        self.start_angle = angle


    def set_end_angle(self, angle):
        """
        Sets the end angle for the program to 'angle', where
        'angle' is in degrees. This method may be called while the 
        scan program is running. However, the end angle listed at 
        the top of the angles file will not be changed.
        """
        self.end_angle = angle
    
    
    def set_duration(self, duration):
        """
        Sets the duration for the program to 'duration', where
        'duration' is in minutes. This method may be called while the 
        scan program is running. However, the scan time listed at 
        the top of the angles file will not be changed.
        """
        self.duration = duration
        vel = abs(self.end_angle - self.start_angle) / self.duration 
        if self._scanner.get_velocity() != vel:
            self._scanner.set_velocity(vel)
        
        
    def _run(self):
        """
        Implements the actual scan program - in this case continuously moves the 
        motor back and forth between the start and end angles at a velocity such 
        that the scan takes the required duration.

        """     
        #set the required velocity
        self.set_duration(self.duration)

        while True:
            self._scanner.set_position(self.end_angle)
            
            if self._exit_flag:
                break
            
            self._scanner.set_position(self.start_angle)
            
            if self._exit_flag:
                break
        
        

class ScanProgramLogged(ScanProgramSimple):
    """
    Scan program that logs the scan angles to a file
    
    * scanner - a scanner.Scanner object
    * start - start angle in degrees
    * end - end angle in degrees
    * duration - scan time (one way) in minutes.
    * filename - the file to log the angles to.
    * extra_procs - extra processes that need to be run at the same time as the scan program
      for example, this is how the file system monitor threads are started and ended. It should
      be a list of objects that implement a start and stop method. The start method will be called
      on scan program start and stop will be called on scan program end. The start method must 
      return immediately.
    """
    
    def __init__(self, scanner, start, end, duration, filename, field_notes, extra_procs=None):
        ScanProgramSimple.__init__(self, scanner, start, end, duration, extra_procs=extra_procs)
        self.filename = filename
        self.notes = field_notes
        self.log_file_lock = threading.Lock()  
    
    
    def _run(self):
        """
        Implements the actual scan program - in this case continuously moves the 
        motor back and forth between the start and end angles at a velocity such 
        that the scan takes the required duration and log the scan positions to file.
        Note, that since the acceleration time of the motor is negligible, only the
        end points of the scan will be logged to the file.

        """     
        #calculate the required velocity
        vel = abs(self.end_angle - self.start_angle) / self.duration 
        self._scanner.set_velocity(vel)
        
        self._write_header_data()
        self._record_current_position()
        while True:
            self._scanner.set_position(self.end_angle)
            
            self._record_current_position()
            if self._exit_flag:
                break
            
            self._scanner.set_position(self.start_angle)
            
            self._record_current_position()
            if self._exit_flag:
                break

    
    def _write_header_data(self):
        """
        Writes the header data to the top of the angles file.
        """
        with self.log_file_lock:
            with open(self.filename, 'w') as ofp:
                #write automatically generated header data to file
                ofp.write("###########   AvoScan scan angles   ###########\n#\n")
                ofp.write("# Date (UTC): " + datetime.datetime.utcnow().strftime("%d/%m/%Y") + "\n")
                ofp.write("# Start time (UTC): " + datetime.datetime.utcnow().strftime("%H:%M:%S") + "\n")
                ofp.write("# Scan Mode: Continuous\n")
                ofp.write("# Start angle: " + str(self.start_angle) + "\n")
                ofp.write("# End angle: " + str(self.end_angle) + "\n")
                ofp.write("# Scan duration in minutes(one direction): " + str(round(self.duration, 3)) + "\n")
                if self._scanner.compass is not None:
                    ofp.write("# Scan orientation: " + self._scanner.compass.get_scan_type() + "\n")
                    ofp.write("# Pitch angle: " + str(self._scanner.compass.get_pitch()) + " degrees\n")
                    ofp.write("# Roll angle: " + str(self._scanner.compass.get_roll()) + " degrees\n")
                ofp.write("# Scanner serial number: " + str(self._scanner.get_serial_num()) + "\n")
                ofp.write("#\n###############################################\n#\n")
                
                #write user-input header data to file
                for line in self.notes:
                    ofp.write("# " + line + "\n")
                                
                ofp.write("#\n###############################################\n")
                ofp.write("# Time\t\t\tAngle\n")

    
    def _record_current_position(self):
        """
        Writes the current position of the scanner to the angles file.
        """
        with self.log_file_lock:
            with open(self.filename, 'a') as ofp:
                ofp.write(str(date2secs(datetime.datetime.utcnow())) + "\t" + str(round(self._scanner.get_position(), 3)) + "\n")


class InterruptableSleeper:
    """
    Class providing a sleep() method which can be interrupted from an other thread by calling the
    stop() method. This is used by the SteppedScanProgramLogged class to ensure that the scan 
    program can be stopped instantly. The start() method must be called after each call to stop() 
    before sleep() is called again - otherwise sleep() will return immediately.
    """
    
    def __init__(self):
        self.__q = Queue.Queue(maxsize=1)
        self.__allow_sleep = True
    
    
    def start(self):
        """
        The start() method must be called after each call to stop() before sleep() is called
        again - otherwise sleep() will return immediately.
        """
        #empty the queue (in case a stop request is unprocessed)
        try:
            self.__q.get_nowait()
        except Queue.Empty:
            pass
        
        self.__allow_sleep = True
    
    
    def stop(self):
        """
        Causes a current call to sleep() to return. Subsequent calls to sleep() will return
        immediately unless start() is called again.
        """       
        self.__allow_sleep = False
        try:
            self.__q.put_nowait(None)
        except Queue.Full:
            #a stop request has already been registered
            pass
    
    
    def sleep(self, time):
        """
        Sleep for 'time' seconds, unless a stop request is received.
        """
        if self.__allow_sleep:
            try:
                self.__q.get(block=True, timeout=time)
            except Queue.Empty:
                pass
        

class SteppedScanProgramLogged(ScanProgramBase):
    """
    Scan program that pauses at set angle steps for a prescribed duration. This is useful if you
    want to record spectra at a particular angle, rather than recording as the scanner is moving.
    This program only steps in one direction - when it reaches the end angle it returns the scanner
    to the start angle and then begins stepping again.
    
    * scanner - a scanner.Scanner object
    * start - start angle in degrees
    * end - end angle in degrees
    * num_steps - number of steps to have between the start and end angles
    * pause_time - the time (in seconds) to spend at each step
    * filename - the file to log the angles to.
    * extra_procs - extra processes that need to be run at the same time as the scan program
      for example, this is how the file system monitor threads are started and ended. It should
      be a list of objects that implement a start and stop method. The start method will be called
      on scan program start and stop will be called on scan program end. The start method must 
      return immediately.
    """
    def __init__(self, scanner, start, end, num_steps, pause_time, filename, field_notes, extra_procs=None):
        self.__sleeper = InterruptableSleeper()
        extra_procs.append(self.__sleeper) #start and stop will now be handled by the base class
        ScanProgramBase.__init__(self, scanner, start, extra_procs=extra_procs)
        self.end_angle = float(end)
        self.num_steps = int(num_steps)
        self.pause_time = float(pause_time)
        self.filename = filename
        self.notes = field_notes
        
        
    
    
    def set_duration(self, duration):
        """
        Sets the duration for the program to 'duration', where
        'duration' is in minutes. This method may be called while the 
        scan program is running. However, the scan time listed at 
        the top of the angles file will not be changed.
        """
        self.duration = duration
        vel = abs(self.end_angle - self.start_angle) / self.duration 
        if self._scanner.get_velocity() != vel:
            #need to record our position before changing velocity - otherwise 
            #how would you interpolate to find the angle for a specific time
            self._record_current_position()
            self._scanner.set_velocity(vel)    

        
    def _run(self):
        """
        Implements the actual scan program - in this case moves the motor between
        the start and end angles in equal steps, pausing at each step for the 
        desired amount of time. The times at the start and end of each step are recorded
        to file.
        """     
        #the motor should already be set to tracking speed, but just in case we set it again
        self._scanner.set_to_tracking_speed()
        
        #start writing the angles file
        self._write_header_data()
        
        #calculate the angle step in degrees - remember that this could be negative!
        angle_step = (self.end_angle - self.start_angle) / float(self.num_steps - 1) # -1 because start angle is one step
        
        
        while not self._exit_flag:
            current_step_num = 1
            self._goto_start_pos()
            while True:
                self._record_current_position()
                self.__sleeper.sleep(self.pause_time)
                self._record_current_position()
                
                if self._exit_flag:
                    break
                
                if current_step_num < self.num_steps:
                    self._scanner.step(angle_step)
                    current_step_num += 1
                else:
                    #reached the end of this scan - return to the start angle
                    break
                
                if self._exit_flag:
                    break


    def _write_header_data(self):
        """
        Writes the header data to the top of the angles file.
        """
        with open(self.filename, 'w') as ofp:
            #write automatically generated header data to file
            ofp.write("###########   AvoScan scan angles   ###########\n#\n")
            ofp.write("# Date (UTC): " + datetime.datetime.utcnow().strftime("%d/%m/%Y") + "\n")
            ofp.write("# Start time (UTC): " + datetime.datetime.utcnow().strftime("%H:%M:%S") + "\n")
            ofp.write("# Scan Mode: Stepped\n")
            ofp.write("# Start angle: " + str(self.start_angle) + "\n")
            ofp.write("# End angle: " + str(self.end_angle) + "\n")
            ofp.write("# Number of steps: " + str(self.num_steps) + "\n")
            ofp.write("# Time at each step (seconds): " + str(round(self.pause_time, 3)) + "\n")
            if self._scanner.compass is not None:
                ofp.write("# Scan orientation: " + self._scanner.compass.get_scan_type() + "\n")
                ofp.write("# Pitch angle: " + str(self._scanner.compass.get_pitch()) + " degrees\n")
                ofp.write("# Roll angle: " + str(self._scanner.compass.get_roll()) + " degrees\n")
            ofp.write("# Scanner serial number: " + str(self._scanner.get_serial_num()) + "\n")
            ofp.write("#\n###############################################\n#\n")
            
            #write user-input header data to file
            for line in self.notes:
                ofp.write("# " + line + "\n")
                            
            ofp.write("#\n###############################################\n")
            ofp.write("# Time\t\t\tAngle\n")

    
    def _record_current_position(self):
        """
        Writes the current position of the scanner to the angles file.
        """
        with open(self.filename, 'a') as ofp:
            ofp.write(str(date2secs(datetime.datetime.utcnow())) + "\t" + str(round(self._scanner.get_position(), 3)) + "\n")
        
